import numpy as np
import random

class NSGA2():
    def __init__(self, pop_size=100, n_var=5, pc=0.9, pm=0.1, eta_c=20, eta_m=20):
        self.pop_size = pop_size
        self.n_var = n_var
        self.pc = pc        # crossover probability
        self.pm = pm        # mutation probability
        self.eta_c = eta_c  # SBX distribution index
        self.eta_m = eta_m  # PM distribution index
        self.mating_pool = []
        self.offspring_pool = []
        self.ranks = []
        self.crowding_distances = []

        assert 0 <= self.pc <= 1
        assert 0 <= self.pm <= 1
        assert self.pop_size % 2 == 0
        
    def initialize(self, prob):
        x = prob.lower + (prob.upper - prob.lower) * np.random.rand(
            self.pop_size, prob.n_var)
        return x
    
    def fitness_assignment(self, population_fx):
        pop_size = len(population_fx)
        n_obj = len(population_fx[0])
        
        self.ranks = np.zeros(pop_size, dtype=int)
        self.crowding_distances = np.zeros(pop_size)
        
        dominated_solutions = [[] for _ in range(pop_size)]
        domination_counts = np.zeros(pop_size, dtype=int)
        
        # non-dominated sorting
        for i in range(pop_size):
            for j in range(pop_size):
                if i != j:
                    # check if i dominates j
                    if np.all(population_fx[i] <= population_fx[j]) and np.any(population_fx[i] < population_fx[j]):
                        dominated_solutions[i].append(j)
                    # check if j dominates i
                    elif np.all(population_fx[j] <= population_fx[i]) and np.any(population_fx[j] < population_fx[i]):
                        domination_counts[i] += 1
        
        # front-1
        current_front = np.where(domination_counts == 0)[0].tolist()
        
        # find all fronts
        fronts = [current_front]
        front_no = 0
        
        while current_front:
            next_front = []
            for i in current_front:
                for j in dominated_solutions[i]:
                    domination_counts[j] -= 1
                    if domination_counts[j] == 0:
                        self.ranks[j] = front_no + 1
                        next_front.append(j)
            
            front_no += 1
            if next_front:
                fronts.append(next_front)
                current_front = next_front
            else:
                break
        
        # find crowding_distances
        for front in fronts:
            if len(front) <= 1:
                for i in front:
                    self.crowding_distances[i] = np.inf
                continue
            
            for obj in range(n_obj):
                # sort by fitness
                sorted_front = sorted(front, key=lambda x: population_fx[x][obj])
                
                # boundary point -> infinity
                self.crowding_distances[sorted_front[0]] = np.inf
                self.crowding_distances[sorted_front[-1]] = np.inf
                
                f_max = population_fx[sorted_front[-1]][obj]
                f_min = population_fx[sorted_front[0]][obj]
                
                if f_max > f_min:
                    for i in range(1, len(sorted_front) - 1):
                        self.crowding_distances[sorted_front[i]] += (
                            population_fx[sorted_front[i+1]][obj] - 
                            population_fx[sorted_front[i-1]][obj]
                        ) / (f_max - f_min)

        
        # normalize crowding_distances to [0,1]
        if np.all(np.isfinite(self.crowding_distances)):
            max_cd = np.max(self.crowding_distances)
            if max_cd > 0:
                self.crowding_distances /= max_cd

        return population_fx
    
    def tournament_selection(self, population_x):
        mating_pool = []
        
        for _ in range(self.pop_size):
            # choose 2 individuals at random
            a, b = np.random.randint(0, self.pop_size, 2)
            
            # compare rank and crowding_distance
            if (self.ranks[a] < self.ranks[b] or 
                (self.ranks[a] == self.ranks[b] and 
                 self.crowding_distances[a] > self.crowding_distances[b])):
                winner = population_x[a]
            else:
                winner = population_x[b]
            
            mating_pool.append(winner)
        
        return mating_pool
        
    def crossover(self, pc):
        offspring_pool = []
        
        # for randomness
        np.random.shuffle(self.mating_pool)
        
        # in pairs
        for i in range(0, self.pop_size, 2):
            parent1 = self.mating_pool[i]
            parent2 = self.mating_pool[min(i+1, self.pop_size-1)]
            
            # init children
            child1, child2 = parent1.copy(), parent2.copy()
            
            # SBX crossover
            if np.random.random() <= pc:
                for j in range(len(parent1)):
                    u = np.random.random()
                    
                    if u <= 0.5:
                        beta = (2 * u) ** (1 / (self.eta_c + 1))
                    else:
                        beta = (1 / (2 * (1 - u))) ** (1 / (self.eta_c + 1))
                    
                    child1[j] = 0.5 * ((1 + beta) * parent1[j] + (1 - beta) * parent2[j])
                    child2[j] = 0.5 * ((1 - beta) * parent1[j] + (1 + beta) * parent2[j])

            # add small change if child equals parent
            if np.allclose(child1, parent1) and np.allclose(child1, parent2):
                child1 += np.random.normal(0, 1e-6, size=child1.shape)
            if np.allclose(child2, parent1) and np.allclose(child2, parent2):
                child2 += np.random.normal(0, 1e-6, size=child2.shape)
            
            offspring_pool.extend([child1, child2])
        
        # keep only pop_size children
        offspring_pool = offspring_pool[:self.pop_size]
        
        for child in offspring_pool:
            np.clip(child, self.lower, self.upper, out=child)
            
        return offspring_pool
        
    def mutation(self, pm):
        for i, individual in enumerate(self.offspring_pool):
            for j in range(len(individual)):
                if np.random.random() <= pm:
                    r = np.random.random()
                    if r < 0.5:
                        delta = (2 * r) ** (1 / (self.eta_m + 1)) - 1
                    else:
                        delta = 1 - (2 * (1 - r)) ** (1 / (self.eta_m + 1))
                    
                    # mutation
                    individual[j] += delta * (self.upper[j] - self.lower[j])
                    
                    individual[j] = np.clip(individual[j], self.lower[j], self.upper[j])

                    # move a bit inside
                    if abs(individual[j] - self.lower[j]) < 1e-6:
                        individual[j] += 1e-3  
                    elif abs(individual[j] - self.upper[j]) < 1e-6:
                        individual[j] -= 1e-3
        
        return self.offspring_pool
        
    def environmental_selection(self, combined_population_list):
        population_x, offspring_x = combined_population_list
        
        # to numpy
        if isinstance(population_x, list):
            population_x = np.array(population_x)
        if isinstance(offspring_x, list):
            offspring_x = np.array(offspring_x)
        
        # combine parents and childs
        combined_population = np.vstack([population_x, offspring_x])
        combined_fitness = self.prob.evaluate(combined_population)
        
        pop_size = len(combined_population)
        n_obj = len(combined_fitness[0])
        
        ranks = np.zeros(pop_size, dtype=int)
        crowding_distances = np.zeros(pop_size)
        
        dominated_solutions = [[] for _ in range(pop_size)]
        domination_counts = np.zeros(pop_size, dtype=int)
        
        # non-dominated sorting
        for i in range(pop_size):
            for j in range(pop_size):
                if i != j:
                    # check if i dominates j
                    if np.all(combined_fitness[i] <= combined_fitness[j]) and np.any(combined_fitness[i] < combined_fitness[j]):
                        dominated_solutions[i].append(j)
                    # check if i dominates j
                    elif np.all(combined_fitness[j] <= combined_fitness[i]) and np.any(combined_fitness[j] < combined_fitness[i]):
                        domination_counts[i] += 1
        
        # front-1
        current_front = np.where(domination_counts == 0)[0].tolist()
        
        # find all fronts
        fronts = [current_front]
        front_no = 0
        
        while current_front:
            next_front = []
            for i in current_front:
                for j in dominated_solutions[i]:
                    domination_counts[j] -= 1
                    if domination_counts[j] == 0:
                        ranks[j] = front_no + 1
                        next_front.append(j)
            
            front_no += 1
            if next_front:
                fronts.append(next_front)
                current_front = next_front
            else:
                break
        
        # find crowding_distances
        for front in fronts:
            if len(front) <= 1:
                for i in front:
                    crowding_distances[i] = np.inf
                continue
            
            for obj in range(n_obj):
                sorted_front = sorted(front, key=lambda x: combined_fitness[x][obj])
                
                crowding_distances[sorted_front[0]] = np.inf
                crowding_distances[sorted_front[-1]] = np.inf
                
                f_max = combined_fitness[sorted_front[-1]][obj]
                f_min = combined_fitness[sorted_front[0]][obj]
                
                if f_max > f_min:
                    for i in range(1, len(sorted_front) - 1):
                        crowding_distances[sorted_front[i]] += (
                            combined_fitness[sorted_front[i+1]][obj] - 
                            combined_fitness[sorted_front[i-1]][obj]
                        ) / (f_max - f_min)
        
        # sort by rank and crowding_distances
        indices = np.lexsort((-crowding_distances, ranks))
        selected_indices = indices[:self.pop_size]
        
        # survivor selection
        selected_population = combined_population[selected_indices]
        selected_fitness = combined_fitness[selected_indices]
        
        return [selected_population, selected_fitness]
    
    def run(self, prob, max_gen, seed = 20250408):
        # send random seed
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)

        # init
        self.prob = prob
        self.lower = np.full(self.n_var, prob.lower)
        self.upper = np.full(self.n_var, prob.upper)
        population_x = self.initialize(prob)
        population_fx = prob.evaluate(population_x)
                
        # evolution loop
        for i in range(1, max_gen):

            self.fitness_assignment(population_fx)
            
            self.mating_pool = self.tournament_selection(population_x)
    
            self.offspring_pool = self.crossover(self.pc)
            
            self.offspring_pool = self.mutation(self.pm)
            
            # evaluate offspring
            offspring_x = np.array(self.offspring_pool)
            offspring_fx = prob.evaluate(offspring_x)

            [population_x, population_fx] = self.environmental_selection([population_x, offspring_x])

            # adaptive mutation based on population diversity
            diversity = np.std(population_x)
            if diversity < 1e-4:
                print("Low diversity: reduce eta_m")
                self.eta_m = max(5, self.eta_m*0.99)
            else:
                pass

            # print progress per 10 gens
            if i % 10 == 0:
                print(f"generation {i}/{max_gen}")
       
        return population_x, population_fx